package pool_ip

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/ThreeDotsLabs/watermill"
	"github.com/ThreeDotsLabs/watermill/message"
	"github.com/ThreeDotsLabs/watermill/message/router/middleware"
	"net"
	"os"
	"proxy_resource_test/bootstrap/models"
	"proxy_resource_test/config"
	"proxy_resource_test/dao"
	"proxy_resource_test/services/mill"
	"proxy_resource_test/services/pool_ip/rsp"
	"proxy_resource_test/utils/logging"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const (
	TaskCountrySuccess = "task:stat:success"      // 各个国家统计的数据
	TaskCountryFail    = "task:stat:fail"         // 各个国家统计的数据
	TasksRunningSet    = "tasks:running"          // 集合：正在运行的 taskID 列表
	ConfigKey          = "task:config:pool"       // 配置写入的哈希键
	ConfigChannel      = "task:config:pool:chg"   // 配置变更发布频道
	ControlChannel     = "task:control:pool"      // 控制命令频道
	StatusHashKey      = "task:stat:pool"         // 任务状态哈希键
	TaskCounter        = "task:stat:pool:count"   // 任务当前计数器
	TaskSuccessCounter = "task:stat:pool:success" // 任务成功计数器
	TaskErrCounter     = "task:stat:pool:error"   // 任务失败计数器
	TaskIpRange        = "task:stat:range"        // ip键， 过期时间最大为30分钟
	TaskDelayCount     = "task:delay:count"       // ip键， ip响应时间计数累计
	Ttl                = 3 * 86400 * time.Second  // 组键共同过期时间
)

// 执行的命令
const (
	CmdStart string = "start"
	CmdPause string = "pause"
	CmdStop  string = "stop"
)

// 用于分配机器编号的租约参数
const (
	MaxMachines = 200
	LeaseTTL    = 30 * time.Second
)

type taskRunner struct {
	cancel  context.CancelFunc
	wg      sync.WaitGroup
	iter    *PopularIter
	sem     chan struct{} // 并发信号量
	Popular string        //指定国家
	cfg     *TaskConfig
}

type TaskConfig = models.TaskConfig

var Manager *ManagerStruct

type ManagerStruct struct {
	mu            sync.Mutex
	Tasks         map[string]*taskRunner
	reconcileTick time.Duration // 拉取周期

	// 租约相关字段
	originalSN int
	instanceID string
	sn         int
	renewOnce  sync.Once
	stopRenew  chan struct{}
}

// 获取ip
func getLocalIPv4() string {
	addrList, err := net.InterfaceAddrs()
	if err != nil {
		fmt.Println("get current host ip err: ", err)
		return ""
	}
	var ip string
	for _, address := range addrList {
		if ipNet, ok := address.(*net.IPNet); ok && !ipNet.IP.IsLoopback() {
			if ipNet.IP.To4() != nil {
				ip = ipNet.IP.String()
				break
			}
		}
	}
	return ip
}

func InitDo() {
	idX, _ := os.Hostname()
	idStr := fmt.Sprintf("%s-%s", idX, getLocalIPv4())
	m := &ManagerStruct{
		Tasks:         make(map[string]*taskRunner),
		reconcileTick: 5 * time.Second,
		instanceID:    idStr,
		stopRenew:     make(chan struct{}),
	}
	Manager = m
	// . **同步**抢号——包初始化时一开始就拿到 SN
	if err := Manager.AcquireSN(context.Background()); err != nil {
		logging.Error("initial AcquireSN failed: %v", err)
		return
	}
	go m.reconcileLoop()
}

// AcquireSN 按序尝试租约分配机器编号
func (m *ManagerStruct) AcquireSN(ctx context.Context) error {
	rdb := dao.GetRedis()
	var machines = config.AppConfig.MaxMachines
	if machines < 1 {
		machines = MaxMachines
	}
	// 1) 如果 originalSN 已经被设置，就先试一次
	if m.originalSN != 0 {
		thisKey := fmt.Sprintf("lease:machine:%d", m.originalSN)
		ok, err := rdb.SetNX(thisKey, m.instanceID, LeaseTTL).Result()
		if err != nil {
			return err
		}
		if ok {
			// 恢复到了原来的 SN
			m.sn = m.originalSN
			m.StartRenew(ctx)
			logging.Info("Re-acquired original SN:", m.sn)
			return nil
		}
		// 如果这个键被别人占了，就进入下面的全量扫描
	}
	// 全量扫描部分
	for {
		for i := 1; i <= machines; i++ {
			key := fmt.Sprintf("lease:machine:%d", i)
			ok, err := rdb.SetNX(key, m.instanceID, LeaseTTL).Result()
			if err != nil {
				return err
			}
			if ok {
				m.sn = i
				m.StartRenew(ctx)
				logging.Info("Acquired machine SN:", i)
				return nil
			}
		}
		// 全部被占，稍后重试
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(200 * time.Millisecond):
		}
	}
}

// StartRenew 定期续租
func (m *ManagerStruct) StartRenew(ctx context.Context) {
	m.renewOnce.Do(func() {
		rdb := dao.GetRedis()
		ticker := time.NewTicker(LeaseTTL / 2)
		go func() {
			defer ticker.Stop()
			for {
				select {
				case <-ticker.C:
					// 续约：如果 key 仍属于自己，刷新 TTL
					// 如果已经不属于自己（TTL 到期被其他实例抢走），就重新 AcquireSN
					key := fmt.Sprintf("lease:machine:%d", m.sn)
					val, err := rdb.Get(key).Result()
					if err != nil {
						if strings.Contains(err.Error(), "redis: nil") {
							logging.Warn("lease key %s expired, reacquiring SN", key)
							_ = m.AcquireSN(ctx)
							return
						}
						logging.Error("redis GET %s error: %v", key, err)
						continue
					}
					if val != m.instanceID {
						logging.Warn("lease key %s stolen by %s, reacquiring SN", key, val)
						_ = m.AcquireSN(ctx)
						return
					}
					// 2) 刷新 TTL
					if err := rdb.PExpire(key, LeaseTTL).Err(); err != nil {
						logging.Error("failed to refresh TTL:", err)
					}
				case <-m.stopRenew:
					return
				case <-ctx.Done():
					return
				}
			}
		}()
	})
}

// ReleaseSN 主动释放租约
func (m *ManagerStruct) ReleaseSN() {
	close(m.stopRenew)
	if m.sn > 0 {
		rdb := dao.GetRedis()
		key := fmt.Sprintf("lease:machine:%d", m.sn)
		rdb.Del(key)
	}
}

// reconcileLoop 定时拉取 Redis Set，自动启动/停止本地任务
func (m *ManagerStruct) reconcileLoop() {
	// 仅当机器获取到编号时，才启动任务
	for m.sn == 0 {
		time.Sleep(100 * time.Millisecond)
		logging.Info("未获取到机器编号，等待中...")
	}
	ticker := time.NewTicker(m.reconcileTick)
	defer ticker.Stop()

	for range ticker.C {
		// 从 Redis 读取当前所有需要运行的 taskID
		ids, err := dao.GetRedis().SMembers(TasksRunningSet).Result()
		if err != nil {
			logging.Error("reconcile SMembers err: %v", err)
			continue
		}
		want := make(map[string]struct{}, len(ids))
		for _, id := range ids {

			popular := dao.GetRedis().HGet(fmt.Sprintf("%s:%s", ConfigKey, id), "popular").Val()

			want[id] = struct{}{}
			// 如果本地未跑，则启动，
			//20250722 新增如果任务设置了机器编号
			cfg, err := InitOrFetchConfig(id)
			if err != nil {
				logging.Error("load config error:", err)
				continue
			}
			// 则根据机器编号范围内启动[min,max],否则启动所有机器编号
			// 判断范围：若配置了范围且当前SN不在[min,max]内，跳过
			if cfg.MachineSnMin > 0 && cfg.MachineSnMax > 0 {
				if m.sn < cfg.MachineSnMin || m.sn > cfg.MachineSnMax {
					logging.Info("skip task %s for machineSN=%d", id, m.sn)
					continue
				}
			}
			m.mu.Lock()
			_, ok := m.Tasks[id]
			m.mu.Unlock()
			if !ok {
				logging.Info("reconcile: start %s", id)
				m.Start(id, popular, cfg)
			}
		}

		// 停止本地多余任务
		m.mu.Lock()
		for id, runner := range m.Tasks {
			if _, keep := want[id]; !keep {
				logging.Info("reconcile: stop %s", id)
				runner.cancel()
				delete(m.Tasks, id)
			}
		}
		m.mu.Unlock()
	}
}

func (m *ManagerStruct) Start(taskID, popular string, cfg *TaskConfig) {
	// 幂等加入 Redis
	dao.GetRedis().SAdd(TasksRunningSet, taskID)
	m.mu.Lock()
	defer m.mu.Unlock()
	m.startLocked(taskID, popular, cfg)
}

func (m *ManagerStruct) startLocked(taskID, popular string, cfg *TaskConfig) {
	if _, exists := m.Tasks[taskID]; exists {
		return
	}
	// 已完成的任务不再重启
	statusKey := fmt.Sprintf("%s:%s", StatusHashKey, taskID)
	if dao.GetRedis().HGet(statusKey, "status").Val() == "completed" {
		logging.Info(fmt.Sprintf("task %s already completed", taskID))
		return
	}
	ctx, cancel := context.WithCancel(context.Background())
	iter := NewPopularIter(cfg.Popular)
	runner := &taskRunner{
		cancel:  cancel,
		iter:    iter,
		Popular: popular,
		sem:     make(chan struct{}, cfg.Concurrency),
		cfg:     cfg,
	}
	// 新建 runner
	cctx, cancel := context.WithCancel(ctx)
	m.Tasks[taskID] = runner
	logging.Info("task--mch-sn", m.instanceID)
	runner.wg.Add(1)
	go func() {
		defer runner.wg.Done()
		runTask(cctx, taskID, runner)
	}()
	go cleanMoreThanThirty(taskID)
	logging.Info("task %s started", taskID)
}

func (m *ManagerStruct) Stop(taskID string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if runner, ok := m.Tasks[taskID]; ok {
		runner.cancel()
		runner.wg.Wait()
		delete(m.Tasks, taskID)
		logging.Info("task %s stopped", taskID)
	} else {
		logging.Info("task %s not found", taskID)
	}
	// Redis 层
	dao.GetRedis().SRem(TasksRunningSet, taskID)
}

func InitOrFetchConfig(taskID string) (*TaskConfig, error) {
	key := fmt.Sprintf("%s:%s", ConfigKey, taskID)
	exists, err := dao.GetRedis().Exists(key).Result()
	if err != nil {
		logging.Error("read file exists err:", err)
		return nil, err
	}
	if exists == 0 {
		return nil, fmt.Errorf("config not exists")
	}
	m, err := dao.GetRedis().HGetAll(key).Result()
	if err != nil {
		return nil, err
	}
	total, _ := strconv.Atoi(m["total"])
	sem, _ := strconv.Atoi(m["concurrency"])
	duration, _ := strconv.Atoi(m["duration"])
	minSn, _ := strconv.Atoi(m["machine_no_min"])
	maxSn, _ := strconv.Atoi(m["machine_no_max"])
	return &TaskConfig{
		Operate:      m["operate"],
		Total:        int32(total),
		Concurrency:  sem,
		TargetURL:    m["target_url"],
		UserTpl:      m["user_tpl"],
		User:         m["user"],
		Password:     m["password"],
		ProxyAddr:    m["proxy_addr"],
		Popular:      m["popular"],
		Duration:     int64(duration),
		MachineSnMin: minSn,
		MachineSnMax: maxSn,
		OperateName:  m["operate_name"],
	}, nil
}

var countryNeedRemap = map[string]string{
	"GB": "UK",
}

// runTask 持续测试每个代理并统计去重返回的 IP
// 20250707 新增popular 字段 是否指定国家/或流行地区或随机
func runTask(ctx context.Context, taskID string, tr *taskRunner) {
	// 加载配置
	if tr.cfg == nil {
		cfg, err := InitOrFetchConfig(taskID)
		if err != nil {
			logging.Error("load config error:", err)
			return
		}
		tr.cfg = cfg
	}
	// 更新状态为 running
	statusKey := fmt.Sprintf("%s:%s", StatusHashKey, taskID)
	dao.GetRedis().HSet(statusKey, "status", "running")
	dao.GetRedis().Expire(statusKey, 86400*time.Second)
	// 并发控制
	var wg sync.WaitGroup

	var executed int32
	// 发起请求
	tRsp := rsp.TaskRspGather{
		ProxyAddr:   tr.cfg.ProxyAddr,
		UserTpl:     tr.cfg.UserTpl,
		User:        tr.cfg.User,
		Password:    tr.cfg.Password,
		TargetStr:   tr.cfg.TargetURL,
		OperateName: tr.cfg.OperateName,
	}
	country := tr.Popular // 指定国家或随机字段标识,random /pop /eu_pop
	if tr.cfg.MachineSnMin > 0 || tr.cfg.MachineSnMax > 0 {
		if Manager.sn < tr.cfg.MachineSnMin || Manager.sn > tr.cfg.MachineSnMax {
			logging.Info("skip task  for machineSN= ", Manager.sn)
			time.Sleep(1 * time.Second)
			return
		}
	}
	for {
		select {
		case <-ctx.Done():
			// 收到停止信号
			wg.Wait()
			dao.GetRedis().HSet(statusKey, "status", "stopped")
			return

		case tr.sem <- struct{}{}:
			// 信号量加1
			counterKey := fmt.Sprintf("%s:%s", TaskCounter, taskID)
			cur, err := dao.GetRedis().Incr(counterKey).Result()
			if err != nil {
				<-tr.sem
				logging.Error("incr count error:", err)
				continue
			}

			dao.GetRedis().Expire(counterKey, Ttl) //设置过期

			nowUnix := time.Now().Unix()
			if tr.cfg.Total > 0 && (cur-1) >= int64(tr.cfg.Total) ||
				(tr.cfg.Duration > 0 && nowUnix >= tr.cfg.Duration) {
				logging.Info("reached total limit or duration, stopping")
				<-tr.sem
				// 结束任务
				pipe := dao.GetRedis().Pipeline()
				pipe.SRem(TasksRunningSet, taskID)
				taskKey := fmt.Sprintf("%s:%s", StatusHashKey, taskID)
				pipe.HSet(taskKey, "status", "completed")
				pipe.Expire(taskKey, Ttl)
				_, _ = pipe.Exec()
				wg.Wait()
				return
			}

			wg.Add(1)
			// 检查 total 限制， 默认最多到100W 次
			//任务设置了时长限制，且到达了超时时间就停止，设置时效（任务个数默认为0 ）
			//任务如果设置了总执行此处， 且到达了总执行次数就停止设置个数超时时间为0
			go func(task rsp.TaskRspGather, country, optName string) {
				var area string
				if country == "pop" || country == "eu_pop" {
					area = tr.iter.Next()
					if optName == "netnut" {
						if rem, ok := countryNeedRemap[area]; ok {
							area = rem
						}
					}
				}
				defer wg.Done()
				defer func() { <-tr.sem }()
				atomic.AddInt32(&executed, 1) // 增加计数器

				_, TimeSince, infoIp := task.PerformRequest(area)

				infoIp.CountryAppoint = area // 指定国家热门

				infoIp.CreatedAt = time.Now().Unix()
				infoIp.Operator = taskID
				infoIp.Delay = TimeSince
				if infoIp.Country == "" {
					infoIp.State = 2
				}
				infoIp.Sn = Manager.sn
				data, err := json.Marshal(infoIp)
				if err != nil {
					logging.Error("JSON 序列化失败:", err)
					return
				}
				msgID := watermill.NewUUID()
				msg := message.NewMessage(msgID, data)
				middleware.SetCorrelationID(msgID, msg)
				if err := mill.PubSub.Publish(mill.IpInfoRegion, msg); err != nil {
					return
				}
			}(tRsp, country, tr.cfg.OperateName)

			fmt.Println("now-nums", atomic.LoadInt32(&executed), "sem", tr.cfg.Concurrency)
		default:
		}
	}
}

func runTaskV2(ctx context.Context, taskID string, tr *taskRunner) {
	rdb := dao.GetRedis()
	statusKey := fmt.Sprintf("%s:%s", StatusHashKey, taskID)

	// 1. 更新状态为 running
	rdb.HSet(statusKey, "status", "running")
	rdb.Expire(statusKey, 86400*time.Second)

	// 2. 计数器键，提前设置过期一次
	counterKey := fmt.Sprintf("%s:%s", TaskCounter, taskID)
	rdb.Expire(counterKey, Ttl)

	// 3. 准备请求模板
	tRsp := rsp.TaskRspGather{
		ProxyAddr: tr.cfg.ProxyAddr,
		UserTpl:   tr.cfg.UserTpl,
		User:      tr.cfg.User,
		Password:  tr.cfg.Password,
		TargetStr: tr.cfg.TargetURL,
	}
	country := tr.Popular

	var wg sync.WaitGroup

loop:
	for {
		// 4. 检查停止信号
		select {
		case <-ctx.Done():
			break loop
		default:
		}

		// 5. 检查机器编号范围
		if tr.cfg.MachineSnMin > 0 || tr.cfg.MachineSnMax > 0 {
			if Manager.sn < tr.cfg.MachineSnMin || Manager.sn > tr.cfg.MachineSnMax {
				logging.Info("skip task %s for machineSN=%d", taskID, Manager.sn)
				time.Sleep(time.Second)
				continue
			}
		}

		// 6. 抢并发位
		select {
		case tr.sem <- struct{}{}:
		case <-ctx.Done():
			break loop
		}

		// 7. 全局自增，获取序号
		cur, err := rdb.Incr(counterKey).Result()
		if err != nil {
			<-tr.sem
			logging.Error("incr count error: %v", err)
			continue
		}

		// 8. 检查总次数和时长限制
		now := time.Now().Unix()
		if (tr.cfg.Total > 0 && cur > int64(tr.cfg.Total)) ||
			(tr.cfg.Duration > 0 && now >= tr.cfg.Duration) {
			<-tr.sem
			break loop
		}

		// 9. 真正发起请求
		wg.Add(1)
		go func(seq int64, area string) {
			defer wg.Done()
			defer func() { <-tr.sem }()

			// 如果是“pop”，动态获取
			if area == "pop" {
				area = tr.iter.Next()
			}
			_, delay, info := tRsp.PerformRequest(area)
			info.CountryAppoint = area
			info.CreatedAt = time.Now().Unix()
			info.Operator = taskID
			info.Delay = delay
			if info.Country == "" {
				info.State = 2
			}
			fmt.Println(fmt.Printf("task %s: seq=%d concurrent=%d", taskID, seq, len(tr.sem)))
		}(cur, country)
	}

	// 10. 等待所有 in-flight 完成
	wg.Wait()

	// 11. 更新最终状态
	finalStatus := "stopped"
	if tr.cfg.Total > 0 {
		lastCount, _ := dao.GetRedis().Get(counterKey).Int64()
		if lastCount >= int64(tr.cfg.Total) {
			finalStatus = "completed"
		}
	}
	rdb.HSet(statusKey, "status", finalStatus)
	rdb.SRem(TasksRunningSet, taskID)
}

// cleanMoreThirty 清理过期数据
func cleanMoreThanThirty(taskID string) {
	ticker := time.NewTicker(time.Minute)
	defer ticker.Stop()
	keys := fmt.Sprintf("%s:%s", TaskIpRange, taskID)
	for now := range ticker.C {
		threshold := now.Unix() - 86420
		if err := dao.GetRedis().ZRemRangeByScore(keys, "0", strconv.FormatInt(threshold, 10)).Err(); err != nil {
			logging.Error("clean expired zset error:", err, "key:", keys)
		}
	}
}

// TaskMonitor 静默任务执行监听器，这是使用redis 的发布订阅部分代码，此处不用，如果使用方式 为 “test:start ...”
func TaskMonitor() {
	pubsub := dao.GetRedis().Subscribe(ConfigChannel, ControlChannel)
	defer pubsub.Close()

	logging.Info("[master] taskMonitor started")
	for msg := range pubsub.Channel() {
		parts := strings.SplitN(msg.Payload, ":", 2)
		if len(parts) != 2 {
			continue
		}
		taskID, cmd := parts[0], parts[1]
		switch cmd {
		case CmdStart:
			cfg, _ := InitOrFetchConfig(taskID)
			Manager.Start(taskID, "pop", cfg)
		case CmdStop, CmdPause:
			Manager.Stop(taskID)
		}
	}
}
